/**
 * Copyright 2024/3/23 the original brandnewdata.com.
 * <p>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * https://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.zlb.envmodify.transform;

import com.alibaba.dcm.DnsCacheManipulator;
import com.zlb.envmodify.service.EnvModifyService;
import com.zlb.envmodify.service.EnvModifyServiceFactory;
import com.zlb.envmodify.transform.config.EnvSettingConfig;
import com.zlb.envmodify.transform.transformer.ClassTransformer;
import com.zlb.envmodify.transform.transformer.EarlyClassCacheTransformer;
import javassist.ClassPool;
import javassist.CtClass;

import java.lang.instrument.Instrumentation;
import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.function.Supplier;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * @Description 1.基于配置文件增加一些环境变量  2.基于配置文件修改域名解析 3.基于增强代理配置对class进行动态修改
 * @ClassName ModifyClass
 * @Author zhulb
 * @Date 2024/3/23 0:41
 * @Version 0.0.1
 */
public class DynamicEnvModify {

    static Logger logger = Logger.getLogger(DynamicEnvModify.class.getName());

    static EnvModifyService envModifyService = EnvModifyServiceFactory.getSigleton();

    static Instrumentation cacheInst;

    public static void modify(String agentArgs, Instrumentation inst) {
        cacheInst = inst;
        //1.换成转换器生效前就加载的类信息
        EarlyClassCacheTransformer earlyClassCacheTransformer = addEarlyClassCacheTransformer(agentArgs, inst);
        //2.初始化EnvModifyService，设置agent参数
        EnvModifyService envModifyService = EnvModifyServiceFactory.getSigleton();
        envModifyService.setAgentArgs(agentArgs);
        logger.info("SYSTEM_CODE=" + envModifyService.getSystemCode());
        //3.添加字节码修改转换器，同时移除用于缓存的转换器
        ClassTransformer classTransformer = new ClassTransformer();
        inst.addTransformer(classTransformer, true);
        inst.removeTransformer(earlyClassCacheTransformer);
        //4.所有提前加载类需要重新转换
        reTransformEarlyClass(inst, classTransformer, earlyClassCacheTransformer);
        //3.替换环境变量信息
        setEnvAndDns();
    }

    private static void reTransformEarlyClass(Instrumentation inst, ClassTransformer classTransformer, EarlyClassCacheTransformer earlyClassCacheTransformer) {
        for (Class cls : inst.getAllLoadedClasses()) {
            if (classTransformer.contains(cls)) {
                reTransformEarlyClass(inst, cls, earlyClassCacheTransformer);
            }
        }
    }

    private static void reTransformEarlyClass(Instrumentation inst, Class cls, EarlyClassCacheTransformer earlyClassCacheTransformer) {
        try {
            byte[] classByte = earlyClassCacheTransformer.getClassByte(cls.getName().replace(".", "/"));
            //早期缓存找不到时，从ClassPool中获取，并添加到早期缓存(因为还需要借此判断是否可以拷贝原始方法)
            if (classByte == null) {
                logger.log(Level.WARNING, cls.getName() + " 类在早期类缓存前进行加载，使用ClassPool获取字节码!");
                CtClass ctClass = ClassPool.getDefault().get(cls.getName());
                classByte = ctClass.toBytecode();
                earlyClassCacheTransformer.put(cls.getName().replace(".", "/"), classByte);
            }
            //1. 通过retransform方式代理
            inst.retransformClasses(cls);
            //2. 通过重新加载类定义的方式代理
            //ClassDefinition classDefinition = new ClassDefinition(cls, classByte);
            //inst.redefineClasses(classDefinition);
        } catch (Exception e) {
            logger.log(Level.WARNING, e, () -> cls.getName() + "类代理失败");
        }
    }

    public static void hotTransform(Class cls){
        try {
            cacheInst.retransformClasses(cls);
        } catch (Exception e){
            logger.log(Level.WARNING, e, () -> cls.getName() + "类代理失败");
        }
    }

    private static EarlyClassCacheTransformer addEarlyClassCacheTransformer(String agentArgs, Instrumentation inst) {
        /**
         * 在我们的转换器生效前，JVM已经加载了很多的类了，比如Properties类
         * 而且在我们new一个转换器的时候，JVM也会加载我们用到的类
         * 这些类都属于提前加载的类，所以不会被转换，因此将类的字节信息缓存起来
         * 在后续做重新加载或者重新定义时以供使用
         */
        EarlyClassCacheTransformer earlyClassCacheTransformer = EarlyClassCacheTransformer.getSington();
        inst.addTransformer(earlyClassCacheTransformer);
        return earlyClassCacheTransformer;
    }

    private static void setEnvAndDns() {
        //1.获取配置环境替换配置信息
        EnvSettingConfig envSettingConfig = envModifyService.getEnvSettingConfig();
        Properties properties = envSettingConfig.getProperties();
        //2.基于环境替换配置信息，替换yaml/yml/properties中的配置项
        setEnvAndDns(() -> {
            return envModifyService.getSystemCode();
        }, properties);

    }

    private static void setEnvAndDns(Supplier<String> getSystemCode, Properties properties) {
        //1.全局配置
        setEnvAndDns("all", properties);
        //2.项目配置
        String systemCode = getSystemCode.get();
        if (systemCode == null || "".equals(systemCode)) {
            logger.warning("systemCode为空！");
        } else {
            logger.info("SYSTEM_CODE=" + systemCode);
        }
        //3.如果systemCode不为空，将加载对应的配置
        Optional.ofNullable(systemCode).ifPresent(code -> setEnvAndDns(systemCode, properties));
    }

    private static void setEnvAndDns(String systemCode, Properties properties) {
        setEnv(systemCode, properties);
        setDns(systemCode, properties);
    }

    private static void setEnv(String systemCode, Properties properties) {
        Set<String> propertyNames = properties.stringPropertyNames();
        propertyNames.stream().filter(q -> q.startsWith(envModifyService.getConfigEnvPrefix() + "." + systemCode)).forEach(e -> {
            String key = e.replace(envModifyService.getConfigEnvPrefix() + "." + systemCode + ".", "");
            String value = properties.getProperty(e);
            System.setProperty(key, value);
        });
    }

    private static void setDns(String systemCode, Properties properties) {
        Set<String> propertyNames = properties.stringPropertyNames();
        propertyNames.stream().filter(q -> q.startsWith(envModifyService.getConfigDnsPrefix() + "." + systemCode)).forEach(e -> {
            String domain = e.replace(envModifyService.getConfigDnsPrefix() + "." + systemCode + ".", "");
            String ip = properties.getProperty(e);
            DnsCacheManipulator.setDnsCache(domain, ip);
        });
    }

}
